package org.example.controller;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentBySentenceSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.store.embedding.*;
import lombok.extern.slf4j.Slf4j;
import org.example.service.Assistant;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import java.util.ArrayList;
import java.util.List;

@RestController
@RequestMapping("/api/chat")
@Slf4j
public class ChatController {

    @Autowired
    ChatLanguageModel chatLanguageModel;

    @Autowired
    Assistant assistant;

    @Autowired
    EmbeddingStore<TextSegment> embeddingStore;

    @Autowired
    EmbeddingModel embeddingModel;


    @GetMapping("/low/hello")
    public String hello(@RequestParam(name = "message") String message) {
        return chatLanguageModel.chat(UserMessage.from(message)).aiMessage().text();
    }

    @GetMapping("/high/hello")
    public String highHello(@RequestParam(name = "message") String message) {
        return assistant.chat(message);
    }

    @GetMapping("/low/role")
    public String lowRoleHello(@RequestParam(name = "message") String message) {
        return chatLanguageModel.chat(List.of(
                SystemMessage.systemMessage("假如你是雷军，接下来你用雷军的语气来对话"),
                UserMessage.from(message))).aiMessage().text();
    }

    @GetMapping("/high/role")
    public String highRoleHello(@RequestParam(name = "message") String message) {
        return assistant.chatWithRole(message);
    }

    @GetMapping("/low/history")
    public String lowHistory(@RequestParam(name = "message") String message) {

        // 最多保存 15 条消息
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(15);
        chatMemory.add(UserMessage.from(message));
        ChatResponse response = chatLanguageModel.chat(chatMemory.messages());
        chatMemory.add(response.aiMessage());
        return response.aiMessage().text();
    }

    @GetMapping("/high/history")
    public String highHistory(@RequestParam(name = "memoryId") String memoryId, @RequestParam(name = "message") String message) {
        return assistant.highHistory(memoryId, message);
    }

    @GetMapping("/load/rag")
    public String load() {
        // 定义分割器
        DocumentBySentenceSplitter splitter = new DocumentBySentenceSplitter(20, 10);
        List<Document> documents = FileSystemDocumentLoader.loadDocuments("G:\\刘续\\学习\\后端\\langchain\\study-langchain\\documents", new TextDocumentParser());
        EmbeddingStoreIngestor.ingest(documents, embeddingStore);

//        // 加载文档
//        List<Document> documents = FileSystemDocumentLoader.loadDocuments("G:\\刘续\\学习\\后端\\langchain\\study-langchain\\documents", new TextDocumentParser());
//        // 切割文档
//        List<TextSegment> textSegments = splitter.splitAll(documents);
//        System.out.println("textSegments.size() = " + textSegments.size());
//        // 转换为向量
//        List<Embedding> embeddings = embeddingModel.embedAll(textSegments).content();
//        System.out.println("embeddings = " + embeddings);
//        EmbeddingStoreIngestor.ingest(documents, embeddingStore);
//        embeddingStore.addAll(embeddings, textSegments);
        return "success";
    }

    @GetMapping("/high/stream")
    public List<String> highStream(@RequestParam(name = "message") String message) {
        List<String> responseList = new ArrayList<>();
        TokenStream tokenStream = assistant.chatStream(message);
        tokenStream.onPartialResponse(item -> {
                    log.info(item);
                    responseList.add(item);
                })
                .onError((Throwable error) -> error.printStackTrace())
                .start();
        return responseList;
    }

    @GetMapping("/embedding/chat")
    public String embeddingChat(@RequestParam(name = "message") String message) {
        Response<Embedding> embed = embeddingModel.embed(message);
        EmbeddingSearchRequest request = EmbeddingSearchRequest.builder()
                .queryEmbedding(embed.content()).build();
        EmbeddingSearchResult<TextSegment> search = embeddingStore.search(request);
        String result = "";
        for (EmbeddingMatch<TextSegment> match : search.matches()) {
            System.out.println(match.embedded().text() + ", 分数为: " + match.score());
//            result = match.embedded().text();
        }
        return result;


    }

}
